
# Set up the packages
library(tidyverse)
library(haven)
set.seed(42)

# Set the path to the working directory
if (Sys.info()[['user']] == "alal") {
  root = "/home/alal/Dropbox/1_Research/ElecAdminFunding/ElecAdminFundingReplication"
} else {
  root = "~/Dropbox/ElecAdminFunding/ElecAdminFundingReplication"
}

##
# Set up the data
##

# Bring in the analysis data
ctcl = read_dta(file.path(root, "modified_data/ctcl_w_small_counties.dta")) %>%
	filter(policy_year==2020) %>%
	select(fips, treat, lag_vs_dem4_pres)

# Run the selection regression
selection_reg = lm(treat ~ lag_vs_dem4_pres, data = ctcl)

# Set up a dataset with electoral votes by state in 2020
elec_votes = read_csv(file.path(root, "/original_data/electoral_votes.csv"))

# Load in the effect estimates
load(file = file.path(root, "tmp/allsynth.RData"), verbose = TRUE)
estimators = c("Time Weighted DID", "DIFP (Regularized)", "Synthetic DID (SDID)")
effect_ests = data.frame(vs_effect_est = allsynth_dem$summary_table[1, estimators], 
  turnout_effect_est = allsynth_tur$summary_table[1, estimators])

# Bring in the presidential election data for simulating a counterfactual
pres = read_dta(file.path(root, "/modified_data/wi_analysis_file_w_missing_pop.dta")) %>%
	filter(year>=2016) %>%
	rename(vap=est_vap) %>%
	select(place, year, votes_dem, votes_rep, vap, treat) %>%
	mutate(state = "WI", fips=550000000 + place) %>%
	pivot_wider(names_from = year,
    values_from = c(votes_dem, votes_rep, vap),
    names_sep = "") %>%
	rename(vap = vap2020) %>%
	select(-vap2016) %>%
	bind_rows(
		read_dta(file.path(root, "/modified_data/pres_results_for_counterfactual.dta"))%>%
		left_join(ctcl, by="fips") %>%
		filter(state!="WI") 
	) %>%
	mutate(votes_dem_rep2020 = votes_dem2020 + votes_rep2020,
		vs_dem2020 = votes_dem2020/votes_dem_rep2020,
		vs_dem2016 = votes_dem2016/(votes_dem2016 + votes_rep2016)) %>%

	# Fill in some gaps in treatment status
	mutate(
		treat = ifelse(state=="CO" & county=="Broomfield", 1, treat),
		treat = ifelse(state=="MD" & county=="Baltimore" & rowtype=="City", 1, treat),
		treat = ifelse(state=="MO" & county=="St. Louis" & rowtype=="City", 1, treat),
		treat = ifelse(state=="NV" & county=="Carson City" & rowtype=="City", 0, treat),
		treat = ifelse(state=="SD" & county=="Oglala Lakota", 0, treat) # managed by Fall River County which did not receive money
	) %>%

	# Impute treatment probabilities based on the selection regression
	mutate(prob_treat = selection_reg$coef[1] + vs_dem2016*selection_reg$coef[2]) 


##
# Run the simulation
##

# Create a dataset for storing the simulated counterfactual
sim_output = tibble(turnout_effect=NA, vs_effect=NA, 
	dem_evotes=NA, 
	.rows=1000*3)

# Simulate 1,000 grant assignments in places with 
# municipal-level election administration
row = 0
for(i in 1:1000){
	if(i%%25==0) print(i)

	# Loop over the three different weighted diff-in-diff estimators
	for(e in 1:3){

		# Update the row number
		row = row + 1

		# Load in the right effect estimate and save it in the output matrix
		turnout_effect = effect_ests$turnout_effect_est[e]/100
		vs_effect = effect_ests$vs_effect_est[e]/100
		sim_output$turnout_effect[row] = turnout_effect
		sim_output$vs_effect[row] = vs_effect

		# Run the simulation
		sim_output$dem_evotes[row] = pres %>%

			# Compute counterfactual turnout, vote share, votes, and Dem votes
			mutate(turnout0 = votes_dem_rep2020/vap - turnout_effect,
				vs0 = vs_dem2020 - vs_effect,
				votes0 = turnout0*vap,
				votes_dem0 = votes0*vs0,
				votes_dem1 = votes_dem_rep2020*vs_dem2020) %>%

			# Randomly construct a treatment regime based on selection probabilities
			mutate(treat_sim = ifelse(!is.na(treat), treat, rbinom(1, 1, prob_treat))) %>%

			# Use the observed votes if this county/municipality is untreated
			# in this run of the simulation
			mutate(votes_counterfactual = ifelse(treat_sim==0, 
					votes_dem_rep2020, votes0),
				votes_dem_counterfactual = ifelse(treat_sim==0,
					votes_dem1, votes_dem0)) %>%

			# Compute the total votes and Dem votes by state
			group_by(state) %>%
			summarize(votes_counterfactual = sum(votes_counterfactual),
				votes_dem_counterfactual = sum(votes_dem_counterfactual)) %>%

			# Determine the winner of each state in this counterfactual
			mutate(dem_win = (votes_dem_counterfactual/votes_counterfactual) > 0.5) %>%

			# Count the number of electoral votes Biden would have received
			left_join(elec_votes, by="state") %>%
			filter(dem_win) %>%
			summarize(evotes = sum(evotes)) %>%
			as.numeric()
	}
}

# Print the simulation output across the three estimators
sim_output %>%
	group_by(turnout_effect, vs_effect, dem_evotes) %>%
	summarize(n = n())
